import torch
from torch import nn
import torch.nn.functional as F
import torch.distributed as dist

"""
这段代码实现了大语言模型（LLM）张量并行（Tensor Parallelism, TP） 中的核心组件 ——并行化线性层，通过将线性层权重按不同维度（输入 / 输出）拆分到多个 GPU，降低单卡内存占用，并通过分布式通信汇总计算结果。代码以 LinearBase 为基类，衍生出 5 种适配不同场景的并行线性层，以下从「基类设计→子类功能→并行逻辑→应用场景」展开分析。
一、核心背景：张量并行中的线性层拆分
线性层的本质是矩阵乘法：y = x @ W^T + b（W 为权重矩阵，形状 [output_size, input_size]）。当 input_size 或 output_size 极大（如百万级）时，W 会占用大量内存（如 W=[1e6, 4096] 约 16GB float32）。
张量并行通过拆分 W 的某一维度（输入或输出），让每个 GPU 仅存储 1/tp_size 的权重，计算后通过分布式通信（如 all_reduce）汇总结果，核心思路如下：

列并行（Column Parallel）：拆分 W 的输出维度（output_size，对应矩阵的 “列” 方向），每个 GPU 计算部分输出。
行并行（Row Parallel）：拆分 W 的输入维度（input_size，对应矩阵的 “行” 方向），每个 GPU 计算部分输入的贡献。

四、核心对比：不同并行线性层的差异
线性层类型	TP 拆分维度	权重形状（单卡）	输入处理方式	输出处理方式	适用场景
ReplicatedLinear	无拆分	[output_size, input_size]	完整输入	完整输出（复制）	输出维度小，避免通信开销
ColumnParallelLinear	0（输出）	[O/tp, I]	完整输入	部分输出（需 all_gather）	输出维度大（如 QKV 生成）
MergedColumnParallel	0（输出）	[sum(O_list)/tp, I]	完整输入	合并部分输出（需拆分）	多个线性层合并（如 Feed-Forward）
QKVParallelLinear	0（输出）	[(Q+2KV)/tp, hidden_size]	完整输入	QKV 部分输出（需拆分）	Transformer 注意力层 QKV 生成
RowParallelLinear	1（输入）	[O, I/tp]	部分输入（需拆分）	完整输出（all_reduce）	输入维度大（如 Feed-Forward 第二层）
五、典型应用：Transformer 层的张量并行组合
在 LLM 的 Transformer 解码器层中，通常组合 QKVParallelLinear（列并行）和 RowParallelLinear（行并行），实现完整的张量并行：

注意力层：用 QKVParallelLinear 生成并行 QKV（列并行，拆输出），通过 FlashAttention 计算注意力后，用 RowParallelLinear 做输出投影（行并行，拆输入）。
Feed-Forward 层：第一层用 ColumnParallelLinear（拆输出，如 4096→16384），激活后第二层用 RowParallelLinear（拆输入，如 16384→4096），中间通过 all_reduce 汇总。
总结
这段代码是工业级 LLM 张量并行的核心实现，通过 5 种并行线性层覆盖了不同场景的需求：

无拆分的 ReplicatedLinear 适合小输出；
列并行系列（Column/MergedColumn/QKV）适合拆输出维度，处理大输出；
行并行 RowParallelLinear 适合拆输入维度，处理大输入。

所有类均通过自定义 weight_loader 实现分片权重加载，配合分布式通信（all_reduce/all_gather）汇总结果，最终在降低单卡内存占用的同时，保证模型计算的正确性。
理解这些组件后，可清晰掌握大模型张量并行的工程化逻辑。
"""
def divide(numerator, denominator):
    assert numerator % denominator == 0
    return numerator // denominator


class LinearBase(nn.Module):
    def __init__(
        self,
        input_size: int,       # 输入特征维度（x 的最后一维）
        output_size: int,      # 输出特征维度（y 的最后一维）
        tp_dim: int | None = None,  # TP 拆分维度（0=输出维度，1=输入维度，None=无拆分）
    ):
        """
        核心作用：定义所有并行线性层的通用属性（输入输出尺寸、分布式参数、拆分维度），统一接口，避免代码冗余。
依赖：需提前初始化 torch.distributed 环境（如 dist.init_process_group(backend="nccl")），否则 dist.get_rank() 报错。
        """
        super().__init__()
        self.input_size = input_size
        self.output_size = output_size
        self.tp_dim = tp_dim  # 标记拆分维度，子类需指定
        # 分布式参数：当前 GPU 序号（tp_rank）、总 GPU 数（tp_size）
        self.tp_rank = dist.get_rank()
        self.tp_size = dist.get_world_size()

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        raise NotImplementedError  # 抽象方法，子类必须实现


class ReplicatedLinear(LinearBase):
    """
    1. ReplicatedLinear：无拆分的复制式线性层
核心特点：权重在所有 GPU 上完整复制，无拆分，适合输出维度较小的场景（无需并行）。
适用场景：输出维度小（如 output_size=1024），权重复制后内存占用可接受（如 [1024, 4096] 约 16MB float32），避免通信开销。
    """
    def __init__(self, input_size: int, output_size: int, bias: bool = False):
        super().__init__(input_size, output_size)  # tp_dim=None（无拆分）
        # 权重完整复制（每个 GPU 都存 [output_size, input_size] 的完整 W）
        self.weight = nn.Parameter(torch.empty(self.output_size, self.input_size))
        self.weight.weight_loader = self.weight_loader  # 加载完整权重
        self.bias = nn.Parameter(torch.empty(self.output_size)) if bias else None
        if self.bias is not None:
            self.bias.weight_loader = self.weight_loader

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
        # 直接复制完整权重（无拆分）
        param.data.copy_(loaded_weight)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 普通线性变换（无并行，所有 GPU 计算相同结果）
        return F.linear(x, self.weight, self.bias)


class ColumnParallelLinear(LinearBase):
    """
    ColumnParallelLinear：列并行线性层（拆分输出维度）
核心逻辑：拆分权重 W 的输出维度（tp_dim=0），每个 GPU 仅存储 output_size//tp_size 的权重，计算部分输出（后续需上层汇总）。
关键细节：
输入 x 无需拆分（每个 GPU 都接收完整输入），计算后输出部分结果（仅对应分片的输出维度）。
完整输出需上层通过 dist.all_gather 拼接（如将 4 个 GPU 的 output_size_per_partition=256 拼接为 1024）。
适用场景：输出维度大（如 output_size=4096），拆分后单卡权重内存降低 tp_size 倍（如 4 卡从 16GB 降至 4GB）。
    """
    def __init__(self, input_size: int, output_size: int, bias: bool = False):
        super().__init__(input_size, output_size, tp_dim=0)  # 按输出维度（0）拆分
        # 计算分片尺寸：输入维度不拆分，输出维度按 tp_size 均分
        self.input_size_per_partition = input_size  # 完整输入维度
        self.output_size_per_partition = divide(output_size, self.tp_size)  # 分片输出维度

        # 权重分片：[output_size_per_partition, input_size]（仅存 1/tp_size 的输出维度）
        self.weight = nn.Parameter(torch.empty(self.output_size_per_partition, self.input_size))
        self.weight.weight_loader = self.weight_loader
        self.bias = nn.Parameter(torch.empty(self.output_size_per_partition)) if bias else None
        if self.bias is not None:
            self.bias.weight_loader = self.weight_loader

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
        # 加载权重分片：按 tp_dim=0（输出维度）切分，取当前 rank 对应的部分
        param_data = param.data
        shard_size = param_data.size(self.tp_dim)  # 分片大小（output_size_per_partition）
        start_idx = self.tp_rank * shard_size  # 当前分片的起始索引
        # narrow：沿 tp_dim 切分权重（不复制数据，高效）
        loaded_weight_shard = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)
        param_data.copy_(loaded_weight_shard)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 输入 x 完整（无拆分），与分片权重相乘，得到部分输出（[..., output_size_per_partition]）
        return F.linear(x, self.weight, self.bias)

class MergedColumnParallelLinear(ColumnParallelLinear):
    """
    MergedColumnParallelLinear：合并输出的列并行线性层
核心特点：将多个独立线性层的输出合并为一个，按列并行拆分，减少计算步骤（如同时计算多个 Feed-Forward 层的输出）。
适用场景：模型中有多个输入相同、输出不同的线性层（如 Transformer 的 Feed-Forward 层中 fc1 和 fc2 可合并），通过合并减少 kernel 调用次数，提升效率。
    """
    def __init__(self, input_size: int, output_sizes: list[int], bias: bool = False):
        self.output_sizes = output_sizes  # 多个独立输出的尺寸（如 [2048, 4096]）
        super().__init__(input_size, sum(output_sizes), bias=bias)  # 总输出尺寸=各输出尺寸之和

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: int):
        # 加载指定输出分片的权重（如 loaded_shard_id=0 对应 output_sizes[0]）
        param_data = param.data
        # 计算当前分片在合并权重中的偏移（如 output_sizes=[2048,4096]，shard_id=1 的偏移=2048//tp_size）
        shard_offset = sum(self.output_sizes[:loaded_shard_id]) // self.tp_size
        # 当前分片的尺寸（如 output_sizes[1]=4096，tp_size=4，shard_size=1024）
        shard_size = self.output_sizes[loaded_shard_id] // self.tp_size
        # 切分当前 GPU 负责的合并权重部分
        param_data_shard = param_data.narrow(self.tp_dim, shard_offset, shard_size)
        # 加载对应输出的权重分片
        loaded_weight_shard = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank]
        param_data_shard.copy_(loaded_weight_shard)


class QKVParallelLinear(ColumnParallelLinear):
    def __init__(
            self,
            hidden_size: int,  # 输入隐藏层维度（如 4096）
            head_size: int,  # 每个注意力头的维度（如 128）
            total_num_heads: int,  # Q 的总头数（如 32）
            total_num_kv_heads: int | None = None,  # KV 的总头数（默认=Q的头数，支持 GQA/MQA）
            bias: bool = False,
    ):
        """
        QKVParallelLinear：专为 QKV 设计的列并行线性层
核心特点：继承 ColumnParallelLinear，专门用于 Transformer 注意力层的 Q（查询）、K（键）、V（值）生成，合并 QKV 的线性层并按列并行拆分。
关键设计：
合并 QKV 的线性层：避免单独创建 3 个线性层，减少计算和内存开销。
支持 GQA/MQA：通过 total_num_kv_heads 控制 KV 头数（如 total_num_heads=32、total_num_kv_heads=8，即 32 个 Q 头共享 8 个 KV 头）。
适用场景：Transformer 注意力层的 QKV 生成，是 LLM 张量并行的核心组件（如 LLaMA、GPT-3 的注意力层）。
        """
        self.head_size = head_size
        self.total_num_heads = total_num_heads
        self.total_num_kv_heads = total_num_kv_heads or total_num_heads  # 默认为 Q 的头数
        tp_size = dist.get_world_size()

        # 按 TP 拆分后的头数（每个 GPU 负责的头数）
        self.num_heads = divide(self.total_num_heads, tp_size)  # Q 的分片头数
        self.num_kv_heads = divide(self.total_num_kv_heads, tp_size)  # KV 的分片头数

        # 总输出尺寸 = Q尺寸 + K尺寸 + V尺寸 = (总Q头数 + 2*总KV头数) * 头维度
        output_size = (self.total_num_heads + 2 * self.total_num_kv_heads) * self.head_size
        super().__init__(hidden_size, output_size, bias)  # 调用 ColumnParallelLinear 初始化

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str):
        param_data = param.data
        assert loaded_shard_id in ["q", "k", "v"]  # 标记加载 Q/K/V 的权重

        # 计算 Q/K/V 在合并权重中的偏移和尺寸
        if loaded_shard_id == "q":
            shard_size = self.num_heads * self.head_size  # Q 的分片尺寸
            shard_offset = 0  # Q 在合并权重的起始位置
        elif loaded_shard_id == "k":
            shard_size = self.num_kv_heads * self.head_size  # K 的分片尺寸
            shard_offset = self.num_heads * self.head_size  # K 在合并权重的起始位置（Q 之后）
        else:  # "v"
            shard_size = self.num_kv_heads * self.head_size  # V 的分片尺寸
            shard_offset = self.num_heads * self.head_size + self.num_kv_heads * self.head_size  # V 在 K 之后

        # 切分当前 GPU 负责的 Q/K/V 权重部分
        param_data_shard = param_data.narrow(self.tp_dim, shard_offset, shard_size)
        # 加载对应 Q/K/V 的权重分片
        loaded_weight_shard = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank]
        param_data_shard.copy_(loaded_weight_shard)

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor, loaded_shard_id: str):
        param_data = param.data
        assert loaded_shard_id in ["q", "k", "v"]
        if loaded_shard_id == "q":
            shard_size = self.num_heads * self.head_size
            shard_offset = 0
        elif loaded_shard_id == "k":
            shard_size = self.num_kv_heads * self.head_size
            shard_offset = self.num_heads * self.head_size
        else:
            shard_size = self.num_kv_heads * self.head_size
            shard_offset = self.num_heads * self.head_size + self.num_kv_heads * self.head_size
        param_data = param_data.narrow(self.tp_dim, shard_offset, shard_size)
        loaded_weight = loaded_weight.chunk(self.tp_size, self.tp_dim)[self.tp_rank]
        param_data.copy_(loaded_weight)


class RowParallelLinear(LinearBase):
    """
    RowParallelLinear：行并行线性层（拆分输入维度）
核心逻辑：拆分权重 W 的输入维度（tp_dim=1），每个 GPU 仅存储 input_size//tp_size 的权重，输入 x 拆分后计算，最终通过 dist.all_reduce 汇总完整输出。
关键细节：
输入 x 需提前按输入维度拆分（如 input_size=4096、tp_size=4，每个 GPU 接收 x 的 1024 维分片）。
输出通过 dist.all_reduce 汇总：每个 GPU 计算的是 “分片输入 × 完整输出权重” 的贡献，加和后得到完整输出（无拆分）。
Bias 仅 rank0 加载：避免多个 GPU 重复加 Bias，汇总后加一次即可保证结果正确。
适用场景：输入维度大（如 input_size=4096），拆分后单卡权重内存降低 tp_size 倍，常与 ColumnParallelLinear
配合使用（如 Transformer 的 Feed-Forward 层：ColumnParallelLinear→激活→RowParallelLinear）。
    """
    def __init__(self, input_size: int, output_size: int, bias: bool = False):
        super().__init__(input_size, output_size, tp_dim=1)  # 按输入维度（1）拆分
        # 计算分片尺寸：输入维度按 tp_size 均分，输出维度不拆分
        self.input_size_per_partition = divide(input_size, self.tp_size)  # 分片输入维度
        self.output_size_per_partition = output_size  # 完整输出维度

        # 权重分片：[output_size, input_size_per_partition]（仅存 1/tp_size 的输入维度）
        self.weight = nn.Parameter(torch.empty(self.output_size, self.input_size_per_partition))
        self.weight.weight_loader = self.weight_loader
        # Bias 仅在 rank0 加载（避免重复加和，汇总后加一次即可）
        self.bias = nn.Parameter(torch.empty(self.output_size)) if (bias and self.tp_rank == 0) else None
        if self.bias is not None:
            self.bias.weight_loader = self.weight_loader

    def weight_loader(self, param: nn.Parameter, loaded_weight: torch.Tensor):
        # 加载权重分片：按 tp_dim=1（输入维度）切分，取当前 rank 对应的部分
        param_data = param.data
        shard_size = param_data.size(self.tp_dim)  # 分片大小（input_size_per_partition）
        start_idx = self.tp_rank * shard_size  # 当前分片的起始索引
        loaded_weight_shard = loaded_weight.narrow(self.tp_dim, start_idx, shard_size)
        param_data.copy_(loaded_weight_shard)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # 1. 输入 x 已按输入维度拆分（上层需确保），与分片权重相乘
        # 2. Bias 仅 rank0 加（其他 rank 为 None，避免重复）
        y = F.linear(x, self.weight, self.bias if self.tp_rank == 0 else None)

        # 3. 汇总所有 GPU 的部分结果（all_reduce 加和，得到完整输出）
        if self.tp_size > 1:
            dist.all_reduce(y, op=dist.ReduceOp.SUM)

        return y
